"""时间序列安全的数据增强工具."""

from __future__ import annotations

import logging
from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
from scipy import stats
from sklearn.preprocessing import StandardScaler

logger = logging.getLogger(__name__)

# 尝试导入 imbalanced-learn，如果未安装则设为 None
try:
    from imblearn.over_sampling import SMOTE
    HAS_IMBLEARN = True
except ImportError:
    HAS_IMBLEARN = False
    logger.warning("imbalanced-learn 未安装，SMOTE 功能将不可用")


# 深度学习模型名称列表（包含 temporal_cnn 及旧名 cnn_1d）
DEEP_LEARNING_MODELS = {"lstm", "gru", "temporal_cnn", "cnn_1d"}

# 传统机器学习模型名称列表
TRADITIONAL_ML_MODELS = {
    "logistic_regression",
    "linear_regression",
    "xgboost",
    "lightgbm",
    "random_forest",
    "svm",
    "decision_tree",
}


def is_deep_learning_model(model_type: str) -> bool:
    """判断模型是否为深度学习模型.

    参数
    ----
    model_type:
        模型类型名称（如 "lstm", "gru", "temporal_cnn"）

    返回
    ----
    bool
        如果是深度学习模型返回 True，否则返回 False
    """
    return model_type.lower() in DEEP_LEARNING_MODELS


def check_class_imbalance(y: pd.Series) -> Dict[str, any]:
    """检查类别不平衡程度.

    参数
    ----
    y:
        标签序列（二分类：1=上涨，0=下跌）

    返回
    ----
    Dict[str, any]
        包含以下字段的字典：
        - is_imbalanced: bool，是否严重不平衡
        - positive_ratio: float，正样本比例
        - negative_ratio: float，负样本比例
        - ratio: float，正负样本比例（较大值/较小值）
    """
    positive_count = int(y.sum())
    negative_count = len(y) - positive_count
    total = len(y)

    if total == 0:
        return {
            "is_imbalanced": False,
            "positive_ratio": 0.0,
            "negative_ratio": 0.0,
            "ratio": 1.0,
        }

    positive_ratio = positive_count / total
    negative_ratio = negative_count / total

    # 计算比例（较大值/较小值）
    if positive_count > 0 and negative_count > 0:
        ratio = max(positive_count, negative_count) / min(positive_count, negative_count)
    else:
        ratio = float("inf")

    # 判断是否严重不平衡：比例 > 7:3 或 < 3:7
    is_imbalanced = ratio > (7.0 / 3.0)

    return {
        "is_imbalanced": is_imbalanced,
        "positive_ratio": positive_ratio,
        "negative_ratio": negative_ratio,
        "ratio": ratio,
        "positive_count": positive_count,
        "negative_count": negative_count,
    }


def should_apply_augmentation(
    X_train: pd.DataFrame,
    y_train: pd.Series,
    model_types: List[str],
) -> Dict[str, any]:
    """判断是否需要数据增强.

    参数
    ----
    X_train:
        训练特征数据
    y_train:
        训练标签数据
    model_types:
        要训练的模型类型列表

    返回
    ----
    Dict[str, any]
        包含以下字段的字典：
        - should_apply: bool，是否应该应用数据增强
        - reasons: List[str]，应用增强的原因列表
        - strategies: List[str]，建议的增强策略列表
    """
    reasons = []
    strategies = []

    # 检查数据量
    data_count = len(X_train)
    if data_count < 500:
        reasons.append(f"数据量不足（{data_count} < 500 条）")
        strategies.append("sliding_window")  # 时间窗口滑动增强
        strategies.append("noise_injection")  # 噪声注入

    # 检查类别不平衡
    imbalance_info = check_class_imbalance(y_train)
    if imbalance_info["is_imbalanced"]:
        reasons.append(
            f"类别严重不平衡（比例 {imbalance_info['ratio']:.2f}:1）"
        )
        strategies.append("class_balance")  # 类别平衡增强

    # 检查是否为深度学习模型
    has_dl_model = any(is_deep_learning_model(mt) for mt in model_types)
    if has_dl_model:
        if data_count < 500:
            reasons.append("深度学习模型且数据量不足")
            strategies.append("sliding_window")  # 时间窗口滑动增强
        # 深度学习模型可以适度使用增强提高泛化能力
        strategies.append("feature_transformation")  # 特征变换增强

    # 判断是否应该应用增强
    should_apply = len(reasons) > 0

    # 如果不满足条件，说明不需要增强
    if not should_apply:
        reasons.append("数据量充足且类别平衡，不需要数据增强")

    return {
        "should_apply": should_apply,
        "reasons": reasons,
        "strategies": list(set(strategies)),  # 去重
        "data_count": data_count,
        "imbalance_info": imbalance_info,
        "has_dl_model": has_dl_model,
    }


def sliding_window_augmentation(
    X: pd.DataFrame,
    y: pd.Series,
    window_size: int = 30,
    step: int = 1,
    max_augmented_samples: Optional[int] = None,
    random_seed: int = 42,
) -> Tuple[pd.DataFrame, pd.Series]:
    """时间窗口滑动增强（用于序列模型）.

    对于序列模型（LSTM/GRU/CNN），使用不同起始点的滑动窗口创建多个训练样本。
    窗口长度固定，滑动步长=1，从不同时间点开始创建序列样本。
    确保每个窗口内的数据保持时间顺序，不跨时间边界。

    参数
    ----
    X:
        特征数据框
    y:
        标签序列
    window_size:
        窗口长度（默认 30 日）
    step:
        滑动步长（默认 1）
    max_augmented_samples:
        最大增强样本数（None 表示不限制，但不超过原始样本数的 3 倍）
    random_seed:
        随机种子（默认 42）

    返回
    ----
    Tuple[pd.DataFrame, pd.Series]
        增强后的特征数据和标签数据
    """
    np.random.seed(random_seed)

    if len(X) < window_size:
        # 数据量不足，无法进行滑动窗口增强
        logger.warning(
            f"数据量 {len(X)} 小于窗口大小 {window_size}，跳过滑动窗口增强"
        )
        return X.copy(), y.copy()

    augmented_X_list = []
    augmented_y_list = []

    # 计算最大增强样本数（不超过原始样本数的 3 倍）
    if max_augmented_samples is None:
        max_augmented_samples = len(X) * 3

    # 从不同起始点创建滑动窗口
    max_start_idx = len(X) - window_size
    start_indices = list(range(0, max_start_idx + 1, step))

    # 如果起始点太多，随机选择一部分
    if len(start_indices) > max_augmented_samples:
        np.random.shuffle(start_indices)
        start_indices = start_indices[:max_augmented_samples]

    for start_idx in start_indices:
        end_idx = start_idx + window_size
        window_X = X.iloc[start_idx:end_idx]
        # 对于序列模型，使用窗口最后一个样本的标签
        window_y = y.iloc[end_idx - 1]

        # 将窗口数据展平为一行（用于非序列模型）
        # 注意：这里假设特征已经是序列形式，如果不是，需要调整
        # 为了简化，我们使用窗口的平均值作为增强样本
        augmented_X_list.append(window_X.mean())
        augmented_y_list.append(window_y)

    if len(augmented_X_list) == 0:
        return X.copy(), y.copy()

    augmented_X = pd.DataFrame(augmented_X_list, index=range(len(X), len(X) + len(augmented_X_list)))
    augmented_y = pd.Series(augmented_y_list, index=augmented_X.index)

    logger.debug(
        f"滑动窗口增强：原始样本 {len(X)} 条，增强后 {len(augmented_X)} 条"
    )

    return augmented_X, augmented_y


def noise_injection(
    X: pd.DataFrame,
    y: pd.Series,
    noise_std_ratio: float = 0.01,
    samples_per_original: int = 2,
    random_seed: int = 42,
) -> Tuple[pd.DataFrame, pd.Series]:
    """小幅噪声注入增强.

    仅对特征值（非原始价格）添加高斯噪声。
    噪声标准差 = noise_std_ratio * 特征标准差（默认 1%）。
    确保增强后的特征值仍在合理范围内。

    参数
    ----
    X:
        特征数据框
    y:
        标签序列
    noise_std_ratio:
        噪声标准差比例（默认 0.01，即 1%）
    samples_per_original:
        每个原始样本生成的增强样本数（默认 2）
    random_seed:
        随机种子（默认 42）

    返回
    ----
    Tuple[pd.DataFrame, pd.Series]
        增强后的特征数据和标签数据
    """
    np.random.seed(random_seed)

    augmented_X_list = []
    augmented_y_list = []

    # 计算每个特征的标准差
    feature_stds = X.std()

    for idx in range(len(X)):
        original_row = X.iloc[idx]
        original_label = y.iloc[idx]

        # 为每个原始样本生成多个增强样本
        for _ in range(samples_per_original):
            # 为每个特征添加高斯噪声
            noise = np.random.normal(0, noise_std_ratio * feature_stds)
            augmented_row = original_row + noise

            # 确保特征值在合理范围内
            # 对于 RSI 等指标（0-100），进行截断
            # 这里假设特征值已经在合理范围内，只做基本的边界检查
            augmented_row = augmented_row.clip(lower=X.min(), upper=X.max())

            augmented_X_list.append(augmented_row)
            augmented_y_list.append(original_label)

    if len(augmented_X_list) == 0:
        return X.copy(), y.copy()

    augmented_X = pd.DataFrame(
        augmented_X_list, index=range(len(X), len(X) + len(augmented_X_list))
    )
    augmented_y = pd.Series(augmented_y_list, index=augmented_X.index)

    logger.debug(
        f"噪声注入增强：原始样本 {len(X)} 条，增强后 {len(augmented_X)} 条"
    )

    return augmented_X, augmented_y


def feature_transformation(
    X: pd.DataFrame,
    y: pd.Series,
    random_seed: int = 42,
) -> Tuple[pd.DataFrame, pd.Series]:
    """特征变换增强.

    通过特征变换创建额外特征，包括：
    - 对数变换：对价格特征使用 log 变换
    - 差分变换：计算一阶、二阶差分特征
    - 标准化变换：使用不同时间窗口的统计量

    这些变换作为额外特征加入，而非替换原始特征。

    参数
    ----
    X:
        特征数据框
    y:
        标签序列
    random_seed:
        随机种子（默认 42，虽然这里不使用随机性）

    返回
    ----
    Tuple[pd.DataFrame, pd.Series]
        增强后的特征数据（包含额外特征）和标签数据
    """
    np.random.seed(random_seed)

    X_augmented = X.copy()

    # 对数变换：对正值特征使用 log 变换
    # 只对数值型特征且值都为正的特征进行对数变换
    for col in X.columns:
        if X[col].dtype in [np.float64, np.int64]:
            if (X[col] > 0).all():
                log_col_name = f"{col}_log"
                X_augmented[log_col_name] = np.log1p(X[col])  # 使用 log1p 避免 log(0)

    # 差分变换：计算一阶、二阶差分
    for col in X.columns:
        if X[col].dtype in [np.float64, np.int64]:
            # 一阶差分
            diff1_col_name = f"{col}_diff1"
            X_augmented[diff1_col_name] = X[col].diff().fillna(0)

            # 二阶差分
            diff2_col_name = f"{col}_diff2"
            X_augmented[diff2_col_name] = X[col].diff().diff().fillna(0)

    # 标准化变换：使用不同时间窗口的统计量
    # 20 日窗口标准化
    window_20 = 20
    if len(X) >= window_20:
        scaler_20 = StandardScaler()
        for col in X.columns:
            if X[col].dtype in [np.float64, np.int64]:
                rolling_mean = X[col].rolling(window=window_20, min_periods=1).mean()
                rolling_std = X[col].rolling(window=window_20, min_periods=1).std()
                # 避免除零
                rolling_std = rolling_std.replace(0, 1)
                normalized_col_name = f"{col}_norm20"
                X_augmented[normalized_col_name] = (X[col] - rolling_mean) / rolling_std

    logger.debug(
        f"特征变换增强：原始特征 {len(X.columns)} 个，增强后 {len(X_augmented.columns)} 个"
    )

    return X_augmented, y.copy()


def class_balance_augmentation(
    X: pd.DataFrame,
    y: pd.Series,
    target_ratio: float = 0.5,
    random_seed: int = 42,
) -> Tuple[pd.DataFrame, pd.Series]:
    """类别平衡增强（使用 SMOTE）.

    使用 SMOTE（Synthetic Minority Oversampling Technique）对少数类样本进行过采样。
    仅对特征空间进行插值，不修改时间序列结构。
    平衡后比例控制在 4:6 到 6:4 之间，避免过度平衡。

    参数
    ----
    X:
        特征数据框
    y:
        标签序列
    target_ratio:
        目标正样本比例（默认 0.5，即 1:1 平衡）
    random_seed:
        随机种子（默认 42）

    返回
    ----
    Tuple[pd.DataFrame, pd.Series]
        增强后的特征数据和标签数据
    """
    if not HAS_IMBLEARN:
        logger.warning("imbalanced-learn 未安装，跳过 SMOTE 增强")
        return X.copy(), y.copy()

    np.random.seed(random_seed)

    # 检查当前类别分布
    imbalance_info = check_class_imbalance(y)
    if not imbalance_info["is_imbalanced"]:
        logger.debug("类别已平衡，跳过 SMOTE 增强")
        return X.copy(), y.copy()

    # 确定目标样本数
    positive_count = imbalance_info["positive_count"]
    negative_count = imbalance_info["negative_count"]

    # 计算目标样本数（平衡到 target_ratio）
    if positive_count < negative_count:
        # 正样本是少数类，需要增加正样本
        target_positive = int(negative_count * target_ratio / (1 - target_ratio))
        target_samples = {1: target_positive, 0: negative_count}
    else:
        # 负样本是少数类，需要增加负样本
        target_negative = int(positive_count * (1 - target_ratio) / target_ratio)
        target_samples = {1: positive_count, 0: target_negative}

    # 使用 SMOTE 进行过采样
    smote = SMOTE(
        sampling_strategy=target_samples,
        random_state=random_seed,
        k_neighbors=min(5, min(positive_count, negative_count) - 1),
    )

    try:
        X_resampled, y_resampled = smote.fit_resample(X, y)
        X_resampled = pd.DataFrame(
            X_resampled, columns=X.columns, index=range(len(X), len(X) + len(X_resampled) - len(X))
        )
        y_resampled = pd.Series(
            y_resampled, index=X_resampled.index
        )

        logger.debug(
            f"SMOTE 增强：原始样本 {len(X)} 条（正样本 {positive_count}，负样本 {negative_count}），"
            f"增强后 {len(X_resampled)} 条"
        )

        return X_resampled, y_resampled
    except Exception as e:
        logger.warning(f"SMOTE 增强失败: {e}，返回原始数据")
        return X.copy(), y.copy()


def validate_augmentation(
    X_original: pd.DataFrame,
    y_original: pd.Series,
    X_augmented: pd.DataFrame,
    y_augmented: pd.Series,
    alpha: float = 0.05,
) -> Dict[str, any]:
    """验证增强质量.

    使用 Kolmogorov-Smirnov 检验验证特征分布相似性。
    验证标签分布合理性。
    记录增强信息（样本数、方法、比例）。

    参数
    ----
    X_original:
        原始特征数据
    y_original:
        原始标签数据
    X_augmented:
        增强后的特征数据
    y_augmented:
        增强后的标签数据
    alpha:
        显著性水平（默认 0.05）

    返回
    ----
    Dict[str, any]
        包含以下字段的字典：
        - is_valid: bool，增强是否有效
        - ks_tests: Dict[str, float]，每个特征的 KS 检验 p 值
        - label_distribution: Dict[str, int]，标签分布
        - warnings: List[str]，警告信息列表
    """
    warnings = []
    ks_tests = {}

    # 对每个特征进行 KS 检验
    common_features = set(X_original.columns) & set(X_augmented.columns)
    for feature in common_features:
        try:
            # 只对原始数据中存在的样本进行检验
            original_values = X_original[feature].values
            # 对于增强数据，只取前 len(original_values) 个样本进行比较
            # 或者比较所有增强样本的分布
            augmented_values = X_augmented[feature].values

            # KS 检验
            statistic, p_value = stats.ks_2samp(original_values, augmented_values)
            ks_tests[feature] = p_value

            if p_value < alpha:
                warnings.append(
                    f"特征 {feature} 的分布可能发生了显著变化（p={p_value:.4f}）"
                )
        except Exception as e:
            logger.warning(f"特征 {feature} 的 KS 检验失败: {e}")

    # 检查标签分布
    original_positive = int(y_original.sum())
    original_negative = len(y_original) - original_positive
    augmented_positive = int(y_augmented.sum())
    augmented_negative = len(y_augmented) - augmented_positive

    label_distribution = {
        "original_positive": original_positive,
        "original_negative": original_negative,
        "augmented_positive": augmented_positive,
        "augmented_negative": augmented_negative,
    }

    # 判断增强是否有效
    # 如果大部分特征的 p 值 > alpha，且标签分布合理，则认为增强有效
    valid_ks_count = sum(1 for p in ks_tests.values() if p > alpha)
    total_ks_count = len(ks_tests)
    is_valid = (
        total_ks_count == 0 or valid_ks_count / total_ks_count >= 0.8
    ) and augmented_positive > 0 and augmented_negative > 0

    return {
        "is_valid": is_valid,
        "ks_tests": ks_tests,
        "label_distribution": label_distribution,
        "warnings": warnings,
        "valid_ks_ratio": valid_ks_count / total_ks_count if total_ks_count > 0 else 1.0,
    }


def apply_data_augmentation(
    X_train: pd.DataFrame,
    y_train: pd.Series,
    model_types: List[str],
    random_seed: int = 42,
    max_augmentation_ratio: float = 3.0,
) -> Tuple[pd.DataFrame, pd.Series, Dict[str, any]]:
    """应用数据增强（主入口函数）.

    根据条件判断选择合适的增强策略。
    应用增强方法。
    执行质量控制。
    返回增强后的数据和报告。

    参数
    ----
    X_train:
        训练特征数据
    y_train:
        训练标签数据
    model_types:
        要训练的模型类型列表
    random_seed:
        随机种子（默认 42）
    max_augmentation_ratio:
        最大增强比例（默认 3.0，即增强样本数不超过原始样本数的 3 倍）

    返回
    ----
    Tuple[pd.DataFrame, pd.Series, Dict[str, any]]
        增强后的特征数据、标签数据和增强信息报告
    """
    np.random.seed(random_seed)

    # 1. 判断是否需要增强
    decision = should_apply_augmentation(X_train, y_train, model_types)

    if not decision["should_apply"]:
        logger.info("数据增强未应用：" + "; ".join(decision["reasons"]))
        return (
            X_train.copy(),
            y_train.copy(),
            {
                "applied": False,
                "reasons": decision["reasons"],
                "original_count": len(X_train),
                "augmented_count": len(X_train),
            },
        )

    logger.info(
        f"开始数据增强：{decision['data_count']} 条样本，"
        f"原因：{'; '.join(decision['reasons'])}，"
        f"策略：{', '.join(decision['strategies'])}"
    )

    # 2. 应用增强策略
    X_augmented = X_train.copy()
    y_augmented = y_train.copy()
    applied_methods = []

    original_count = len(X_train)
    max_augmented_count = int(original_count * max_augmentation_ratio)

    # 按策略顺序应用增强
    for strategy in decision["strategies"]:
        if len(X_augmented) >= max_augmented_count:
            logger.info(f"已达到最大增强样本数 {max_augmented_count}，停止增强")
            break

        try:
            if strategy == "sliding_window":
                # 时间窗口滑动增强（仅用于深度学习模型）
                has_dl = any(is_deep_learning_model(mt) for mt in model_types)
                if has_dl:
                    X_window, y_window = sliding_window_augmentation(
                        X_train, y_train, random_seed=random_seed
                    )
                    # 合并增强样本
                    X_augmented = pd.concat([X_augmented, X_window], ignore_index=True)
                    y_augmented = pd.concat([y_augmented, y_window], ignore_index=True)
                    applied_methods.append("sliding_window")

            elif strategy == "noise_injection":
                # 噪声注入增强
                X_noise, y_noise = noise_injection(
                    X_train, y_train, random_seed=random_seed
                )
                # 合并增强样本
                X_augmented = pd.concat([X_augmented, X_noise], ignore_index=True)
                y_augmented = pd.concat([y_augmented, y_noise], ignore_index=True)
                applied_methods.append("noise_injection")

            elif strategy == "feature_transformation":
                # 特征变换增强（直接修改特征，不增加样本数）
                X_augmented, y_augmented = feature_transformation(
                    X_augmented, y_augmented, random_seed=random_seed
                )
                applied_methods.append("feature_transformation")

            elif strategy == "class_balance":
                # 类别平衡增强（SMOTE）
                X_augmented, y_augmented = class_balance_augmentation(
                    X_augmented, y_augmented, random_seed=random_seed
                )
                applied_methods.append("class_balance")

        except Exception as e:
            logger.warning(f"增强策略 {strategy} 应用失败: {e}，继续其他策略")

    # 3. 质量控制
    validation_result = validate_augmentation(
        X_train, y_train, X_augmented, y_augmented
    )

    # 4. 生成报告
    augmentation_info = {
        "applied": True,
        "methods": applied_methods,
        "original_count": original_count,
        "augmented_count": len(X_augmented),
        "augmentation_ratio": len(X_augmented) / original_count if original_count > 0 else 0.0,
        "reasons": decision["reasons"],
        "strategies": decision["strategies"],
        "validation": validation_result,
    }

    # 如果增强后样本数超过限制，进行截断
    if len(X_augmented) > max_augmented_count:
        logger.warning(
            f"增强样本数 {len(X_augmented)} 超过限制 {max_augmented_count}，进行截断"
        )
        X_augmented = X_augmented.iloc[:max_augmented_count]
        y_augmented = y_augmented.iloc[:max_augmented_count]
        augmentation_info["augmented_count"] = len(X_augmented)
        augmentation_info["augmentation_ratio"] = len(X_augmented) / original_count

    logger.info(
        f"数据增强完成：原始 {original_count} 条，增强后 {len(X_augmented)} 条，"
        f"方法：{', '.join(applied_methods)}"
    )

    return X_augmented, y_augmented, augmentation_info

